import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'

import keras
import keras.backend as K
from keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from keras.utils import multi_gpu_model

from config import patience, batch_size, epochs, num_train_samples, num_valid_samples
from data_generator import train_gen, valid_gen, GetRandomSamples
from migrate import migrate_model
from model import build_encoder_decoder, build_refinement
from utils import L1_loss, L2_loss, edge_loss, get_available_cpus, get_available_gpus
from config import img_cols, img_rows
from config import checkpoint_models_path, train_log_path
from data_generator import denormalize_input, denormalize_output

import numpy as np
import cv2

import tensorflow as tf


class MyCbk(keras.callbacks.Callback):
    def __init__(self, model, stage):
        keras.callbacks.Callback.__init__(self)
        self.model_to_save = model
        self.epoc = 0
        self.stage = stage

    def on_epoch_end(self, epoch, logs=None):
        self.epoc += 1
        fmt = checkpoint_models_path + '%s.%08d-%.4f.hdf5'
        self.model_to_save.save(fmt % (self.stage, epoch, logs['val_loss']))

    def on_batch_end(self, batch, logs=None):
        if batch % 100 != 1:
           return
        x, gt, id = GetRandomSamples('train', 1)
        y = self.model_to_save.predict(x)
        x = denormalize_input(x[0])
        y = denormalize_output(y[0, :, :, 0])
        gt = denormalize_output(gt[0, :, :, 0])
        rgb = x[:, :, :3]
        #depth = np.stack([x[:, :, 3], x[:, :, 3], x[:, :, 3]], axis=-1)
        pred = np.stack([y, y, y], axis=-1)
        gt = np.stack([gt, gt, gt], axis=-1)
        #merged = np.concatenate((rgb, depth, pred, gt), axis=1)
        merged = np.concatenate((rgb, pred, gt), axis=1)
        cv2.imwrite(train_log_path + '/%04d_%04d_%s.PNG' % (self.epoc, batch, id[0]), merged)

# predefine stages
STAGE_COARSE = 'coarse'
STAGE_REFINE = 'refine'
STAGE_OVERALL = 'overall'

def train(stage):
    pretrained_path = None
    pretrained_path = checkpoint_models_path + 'coarse.00000079-0.0050.hdf5'

    tensor_board = keras.callbacks.TensorBoard(log_dir=train_log_path,
                                               histogram_freq=0,
                                               write_graph=True,
                                               write_images=True)

    #model_names = checkpoint_models_path + 'portrait.{epoch:08d}-{val_loss:.4f}.hdf5'
    #model_checkpoint = ModelCheckpoint(model_names, monitor='val_loss', verbose=1, save_best_only=True)

    early_stop = EarlyStopping('val_loss', patience=patience)
    reduce_lr = ReduceLROnPlateau('val_loss', factor=0.1, patience=int(patience / 4), verbose=1)

    model = None
    num_gpu = len(get_available_gpus())
    if num_gpu >= 2:
        print('Training on multi-GPUs!!!')
        with tf.device("/cpu:0"):
            if stage == STAGE_COARSE:
                model = build_encoder_decoder(img_rows, img_cols, train=True)
            elif stage == STAGE_REFINE:
                model = build_encoder_decoder(img_rows, img_cols, train=False)
                # freeze the coarse layers
                model.trainable = False
                for _L in model.layers:
                    _L.trainable = False
                # build refine layers
                model = build_refinement(model, train=True)
            elif stage == STAGE_OVERALL:
                model = build_encoder_decoder(img_rows, img_cols, train=True)
                model = build_refinement(model, train=True)
            else:
                assert False
            if pretrained_path is not None:
                model.load_weights(pretrained_path, by_name=True)
            else:
                migrate_model(model)
        local_model = multi_gpu_model(model, gpus=num_gpu)
        # saving the original model and not the multi-gpu model.
        model_checkpoint = MyCbk(model, stage)
    else:
        print('Training on single GPU or on CPU.')
        if stage == STAGE_COARSE:
            model = build_encoder_decoder(img_rows, img_cols, train=True)
        elif stage == STAGE_REFINE:
            model = build_encoder_decoder(img_rows, img_cols, train=False)
            # freeze the coarse layers
            model.trainable = False
            for _L in model.layers:
                _L.trainable = False
            # build refine layers
            model = build_refinement(model, train=True)
        elif stage == STAGE_OVERALL:
            model = build_encoder_decoder(img_rows, img_cols, train=True)
            model = build_refinement(model, train=True)
        else:
            assert False
        if pretrained_path is not None:
            model.load_weights(pretrained_path, by_name=True)
        else:
            migrate_model(model)
        local_model = model
        model_checkpoint = MyCbk(model, stage)

    target = tf.placeholder(dtype='float32', shape=(None, None, None, None))
    if stage == STAGE_COARSE:
        local_model.compile(optimizer='nadam', loss=L2_loss, target_tensors=[target])
    elif stage == STAGE_REFINE:
        local_model.compile(optimizer='nadam', loss=L1_loss, target_tensors=[target])
    elif stage == STAGE_OVERALL:
        local_model.compile(optimizer='nadam', loss=L1_loss, target_tensors=[target])
    local_model.summary()

    # Final callbacks
    callbacks = [tensor_board, model_checkpoint, early_stop, reduce_lr]

    # Start Fine-tuning
    local_model.fit_generator(train_gen(),
                        steps_per_epoch=num_train_samples // batch_size,
                        validation_data=valid_gen(),
                        validation_steps=num_valid_samples // batch_size,
                        epochs=epochs,
                        verbose=1,
                        callbacks=callbacks,
                        use_multiprocessing=True,
                        workers=2)


if __name__ == '__main__':
    train(STAGE_COARSE)